#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
大模型总结器
使用大模型对文本进行总结
"""

import time
from typing import List, Dict, Optional, Any, Tuple
from pathlib import Path
from ..utils import ToolbeltLogger

# 延迟导入，避免在没有安装依赖时出错
try:
    from openai import OpenAI
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False
    OpenAI = None  # type: ignore


class LLMSummarizer:
    """大模型总结器"""
    
    def __init__(self, api_key: Optional[str] = None, model: str = "gpt-3.5-turbo",
                 api_base: Optional[str] = None, max_retries: int = 3,
                 print_output: bool = False):
        """
        初始化大模型总结器
        
        Args:
            api_key (str, optional): OpenAI API密钥
            model (str): 使用的大模型名称
            api_base (str, optional): API基础URL
            max_retries (int): 最大重试次数
            print_output (bool): 是否将模型输出打印到控制台
        """
        self.model = model
        self.max_retries = max_retries
        self.print_output = print_output
        self.client: Any = None
        
        # 如果依赖可用，则初始化OpenAI客户端
        if OPENAI_AVAILABLE and OpenAI is not None:
            # 如果指定了API基础URL，则使用它
            if api_base:
                self.client = OpenAI(api_key=api_key, base_url=api_base)
            else:
                # 如果没有指定API基础URL但提供了API密钥，则使用默认的OpenAI API
                if api_key:
                    self.client = OpenAI(api_key=api_key)
                # 如果是Ollama或没有API密钥，则使用默认的本地Ollama地址
                else:
                    self.client = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
        
        # 如果没有API密钥或依赖不可用，将使用模拟模式
        self.simulation_mode = self.client is None or not OPENAI_AVAILABLE
    
    def summarize_text(self, text: str, filename: str) -> Tuple[str, Optional[Dict[str, int]]]:
        """
        使用大模型对文本进行总结
        
        Args:
            text (str): 要总结的文本
            filename (str): PDF文件名
            
        Returns:
            Tuple[str, Optional[Dict[str, int]]]: 总结内容和token消耗信息
        """
        if self.simulation_mode:
            # 模拟模式，返回示例总结
            summary = f"这是 {filename} 的模拟总结内容。在实际使用中，这里会是通过大模型生成的文献总结。这是一份详细的2000字左右的文献总结，包含了研究背景、方法、主要发现、结论以及未来研究方向等详细内容。"
            
            # 如果需要打印输出，则打印到控制台
            if self.print_output:
                ToolbeltLogger.info(f"\n=== {filename} 的模型输出 ===")
                ToolbeltLogger.info(summary)
                ToolbeltLogger.info("=" * 50)
            
            return summary, None
        
        # 重试机制
        for attempt in range(self.max_retries):
            try:
                # 构造更加专业的提示词 - 学术文献深度解读
                # 对于128k模型，我们不再截断文本，让Ollama自己处理
                prompt = f"""
# 1.0.0.20250910
你是一位精通生物信息学、分子生物学及学术写作的顶级研究助理。你的任务是深度分析以下英文文献，并生成一份严格结构化、便于后续整合至综述文章的中文总结报告。

---

#### 核心指令：

1.  **最高优先级**：**严禁**输出任何与以下内容无关的文本：
    * 任何原文中的章节标题（如 `ABSTRACT`, `INTRODUCTION`, `MATERIALS AND METHODS` 等）。
    * 任何页码、页眉、页脚信息或版权声明。
    * 任何参考文献列表，或与引用格式相关的文本。
    * 任何“本文献总结”或类似的元信息。
2.  **输出格式**：你的输出**必须且仅能**包含以下框架中的内容，且每个部分都必须以中文形式输出。
3.  **信息提炼**：只提炼对综述有价值的核心信息。如果一个部分（如方法）只被一笔带过，请在总结中也简略处理。

---

#### 总结框架（请严格遵循并以中文输出）：

**【文献标题和基本信息】**
- **标题**：提取文献的完整标题。
- **作者**：列出主要作者和通讯作者。
- **机构**：列出主要研究机构。
- **关键词**：提取文献中的核心关键词。

**【1. 核心研究问题与创新点】**
- 简明扼要地阐述本文献试图解决的**核心科学问题**。
- 总结其提出的**核心假设**和**主要创新点**。

**【2. 方法与技术细节】**
- 详细列出本研究使用的**生物信息学工具、算法或实验技术**。例如：使用的软件、数据库、分析流程（如RNA-seq、单细胞测序数据分析）。
- 阐述这些方法如何应用于核心问题。

**【3. 关键发现与数据】**
- **核心发现**：以清晰的条目形式，列出所有重要的研究结果。
- **关键数据**：**直接提取**或**高度概括**原文中支持核心发现的关键数据（如p值、Fold Change、R方值、准确率等）。这些数据应能直接在综述中作为论据。

**【4. 讨论与局限性】**
- **主要讨论**：本文献的发现如何与其他研究（如果原文有提及）建立联系或形成对比？
- **局限性**：明确指出本文献作者提出的**研究不足或技术局限**。

**【5. 未来展望**
- 总结本文献作者提出的**未来研究方向或建议**。

---

**文献原文**：
{text}  # 不再截断文本，让Ollama自己处理
                """
                
                # 调用大模型API
                if self.client is not None:
                    response = self.client.chat.completions.create(
                        model=self.model,
                        messages=[
                            {"role": "system", "content": "你是一位精通生物信息学、分子生物学及学术写作的顶级研究助理。你的任务是深度分析英文文献，并生成一份严格结构化、便于后续整合至综述文章的中文总结报告。请务必以中文形式输出。"},
                            {"role": "user", "content": prompt}
                        ],
                        max_tokens=8192,  # 增加max_tokens以支持更长的输出(4k-8k字)
                        temperature=0.3
                    )
                    
                    summary = response.choices[0].message.content.strip()
                    
                    # 获取token消耗信息
                    token_usage = None
                    if hasattr(response, 'usage') and response.usage is not None:
                        token_usage = {
                            "prompt_tokens": response.usage.prompt_tokens,
                            "completion_tokens": response.usage.completion_tokens,
                            "total_tokens": response.usage.total_tokens
                        }
                    
                    # 如果需要打印输出，则打印到控制台
                    if self.print_output:
                        ToolbeltLogger.info(f"\n=== {filename} 的模型输出 ===")
                        ToolbeltLogger.info(summary)
                        if token_usage:
                            ToolbeltLogger.info(f"\nToken消耗: 提示{token_usage['prompt_tokens']}, 完成{token_usage['completion_tokens']}, 总计{token_usage['total_tokens']}")
                        ToolbeltLogger.info("=" * 50)
                    
                    return summary, token_usage
            except Exception as e:
                ToolbeltLogger.warning(f"第{attempt + 1}次尝试总结 {filename} 失败: {e}")
                if attempt < self.max_retries - 1:
                    ToolbeltLogger.info(f"  {2 ** attempt}秒后进行第{attempt + 2}次重试...")
                    time.sleep(2 ** attempt)  # 指数退避
                else:
                    ToolbeltLogger.warning(f"  已达到最大重试次数，无法生成 {filename} 的总结内容。")
        
        error_summary = f"无法生成 {filename} 的总结内容（已重试{self.max_retries}次）。"
        
        # 即使出错也打印错误信息（如果启用了打印输出）
        if self.print_output:
            ToolbeltLogger.info(f"\n=== {filename} 的模型输出 ===")
            ToolbeltLogger.info(error_summary)
            ToolbeltLogger.info("=" * 50)
        
        return error_summary, None